package snowpipestreaming

import (
	"bufio"
	"context"
	stdjson "encoding/json"
	"errors"
	"fmt"
	"math"
	"net/http"
	"os"
	"slices"
	"strconv"
	"sync"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/samber/lo"
	"github.com/tidwall/gjson"

	"github.com/rudderlabs/rudder-go-kit/bytesize"
	"github.com/rudderlabs/rudder-go-kit/config"
	"github.com/rudderlabs/rudder-go-kit/logger"
	"github.com/rudderlabs/rudder-go-kit/stats"
	"github.com/rudderlabs/rudder-go-kit/stringify"
	obskit "github.com/rudderlabs/rudder-observability-kit/go/labels"

	"github.com/rudderlabs/rudder-go-kit/jsonrs"
	backendconfig "github.com/rudderlabs/rudder-server/backend-config"
	"github.com/rudderlabs/rudder-server/jobsdb"
	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/common"
	snowpipeapi "github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/snowpipestreaming/internal/api"
	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/snowpipestreaming/internal/model"
	"github.com/rudderlabs/rudder-server/utils/misc"
	"github.com/rudderlabs/rudder-server/utils/timeutil"
	"github.com/rudderlabs/rudder-server/warehouse/integrations/manager"
	whutils "github.com/rudderlabs/rudder-server/warehouse/utils"
	"github.com/rudderlabs/rudder-server/warehouse/validations"
)

func New(
	conf *config.Config,
	log logger.Logger,
	statsFactory stats.Stats,
	destination *backendconfig.DestinationT,
) *Manager {
	m := &Manager{
		appConfig: conf,
		logger: log.Child("snowpipestreaming").Withn(
			obskit.WorkspaceID(destination.WorkspaceID),
			obskit.DestinationID(destination.ID),
			obskit.DestinationType(destination.DestinationDefinition.Name),
		),
		statsFactory:        statsFactory,
		destination:         destination,
		now:                 timeutil.Now,
		channelCache:        sync.Map{},
		polledImportInfoMap: make(map[string]*importInfo),
		validator:           validations.NewDestinationValidator(),
	}

	m.config.client.url = conf.GetString("SnowpipeStreaming.Client.URL", "http://localhost:9078")
	m.config.client.maxHTTPConnections = conf.GetInt("SnowpipeStreaming.Client.maxHTTPConnections", 10)
	m.config.client.maxHTTPIdleConnections = conf.GetInt("SnowpipeStreaming.Client.maxHTTPIdleConnections", 5)
	m.config.client.maxIdleConnDuration = conf.GetDuration("SnowpipeStreaming.Client.maxIdleConnDuration", 30, time.Second)
	m.config.client.disableKeepAlives = conf.GetBool("SnowpipeStreaming.Client.disableKeepAlives", true)
	m.config.client.timeoutDuration = conf.GetDuration("SnowpipeStreaming.Client.timeout", 300, time.Second)
	m.config.client.retryWaitMin = conf.GetDuration("SnowpipeStreaming.Client.retryWaitMin", 100, time.Millisecond)
	m.config.client.retryWaitMax = conf.GetDuration("SnowpipeStreaming.Client.retryWaitMax", 10, time.Second)
	m.config.client.retryMax = conf.GetInt("SnowpipeStreaming.Client.retryMax", 5)
	m.config.instanceID = conf.GetString("INSTANCE_ID", "1")
	m.config.maxBufferCapacity = conf.GetReloadableInt64Var(512*bytesize.KB, bytesize.B, "SnowpipeStreaming.maxBufferCapacity")
	m.config.backoff.initialInterval = conf.GetReloadableDurationVar(1, time.Second, "SnowpipeStreaming.backoffInitialIntervalInSeconds")
	m.config.backoff.multiplier = conf.GetReloadableFloat64Var(2.0, "SnowpipeStreaming.backoffMultiplier")
	m.config.backoff.maxInterval = conf.GetReloadableDurationVar(1, time.Hour, "SnowpipeStreaming.backoffMaxIntervalInHours")
	m.config.stuckPipelineThreshold = conf.GetReloadableDurationVar(15, time.Minute, "SnowpipeStreaming.stuckPipelineThresholdInMinutes")

	tags := stats.Tags{
		"module":        "batch_router",
		"workspaceId":   destination.WorkspaceID,
		"destType":      destination.DestinationDefinition.Name,
		"destinationId": destination.ID,
	}
	m.stats.jobs.importing = statsFactory.NewTaggedStat("snowpipe_streaming_jobs", stats.CountType, lo.Assign(tags, stats.Tags{
		"status": "importing",
	}))
	m.stats.jobs.succeeded = statsFactory.NewTaggedStat("snowpipe_streaming_jobs", stats.CountType, lo.Assign(tags, stats.Tags{
		"status": "succeeded",
	}))
	m.stats.jobs.failed = statsFactory.NewTaggedStat("snowpipe_streaming_jobs", stats.CountType, lo.Assign(tags, stats.Tags{
		"status": "failed",
	}))
	m.stats.jobs.aborted = statsFactory.NewTaggedStat("snowpipe_streaming_jobs", stats.CountType, lo.Assign(tags, stats.Tags{
		"status": "aborted",
	}))

	m.stats.discards = statsFactory.NewTaggedStat("snowpipe_streaming_discards", stats.CountType, tags)
	m.stats.pollingInProgress = statsFactory.NewTaggedStat("snowpipe_streaming_polling_in_progress", stats.CountType, tags)

	if m.requestDoer == nil {
		m.requestDoer = m.retryableClient().StandardClient()
	}

	m.api = newApiAdapter(
		m.logger,
		statsFactory,
		snowpipeapi.New(m.appConfig, m.statsFactory, m.config.client.url, m.requestDoer),
		destination,
	)
	m.managerCreator = func(ctx context.Context, modelWarehouse whutils.ModelWarehouse, conf *config.Config, logger logger.Logger, statsFactory stats.Stats) (manager.Manager, error) {
		sf, err := manager.New(whutils.SnowpipeStreaming, conf, logger, statsFactory)
		if err != nil {
			return nil, fmt.Errorf("creating snowflake manager: %w", err)
		}
		err = sf.Setup(ctx, modelWarehouse, whutils.NewNoOpUploader())
		if err != nil {
			return nil, fmt.Errorf("setting up snowflake manager: %w", err)
		}
		return sf, nil
	}
	return m
}

func (m *Manager) retryableClient() *retryablehttp.Client {
	client := retryablehttp.NewClient()
	client.HTTPClient = &http.Client{
		Transport: &http.Transport{
			DisableKeepAlives:   m.config.client.disableKeepAlives,
			MaxConnsPerHost:     m.config.client.maxHTTPConnections,
			MaxIdleConnsPerHost: m.config.client.maxHTTPIdleConnections,
			IdleConnTimeout:     m.config.client.maxIdleConnDuration,
		},
		Timeout: m.config.client.timeoutDuration,
	}
	client.Logger = nil
	client.RetryWaitMin = m.config.client.retryWaitMin
	client.RetryWaitMax = m.config.client.retryWaitMax
	client.RetryMax = m.config.client.retryMax
	return client
}

func (m *Manager) validateConfig(ctx context.Context, dest *backendconfig.DestinationT) error {
	response := m.validator.Validate(ctx, dest)
	if response.Success {
		return nil
	}
	return errors.New(response.Error)
}

func (m *Manager) Now() time.Time {
	return m.now()
}

func (m *Manager) Transform(job *jobsdb.JobT) (string, error) {
	return common.GetMarshalledData(string(job.EventPayload), job.JobID)
}

// Upload uploads data to the Snowpipe streaming destination.
// It reads events from the file, groups them by table, and sends them to Snowpipe.
// It returns the IDs of the importing and failed jobs.
// In case of failure, it aborts the jobs and returns the aborted job IDs.
func (m *Manager) Upload(asyncDest *common.AsyncDestinationStruct) common.AsyncUploadOutput {
	m.logger.Infon("Uploading data to snowpipe streaming destination")

	var destConf destConfig
	err := destConf.Decode(asyncDest.Destination.Config)
	if err != nil {
		return m.abortJobs(asyncDest, fmt.Errorf("failed to decode destination config: %w", err).Error())
	}

	events, err := m.eventsFromFile(asyncDest.FileName, asyncDest.Count)
	if err != nil {
		return m.abortJobs(asyncDest, fmt.Errorf("failed to read events from file: %w", err).Error())
	}
	m.logger.Infon("Read events from file",
		logger.NewIntField("events", int64(len(events))),
		logger.NewIntField("size", int64(asyncDest.Size)),
	)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	discardsChannel, err := m.initializeChannelWithSchema(ctx, asyncDest.Destination.ID, &destConf, discardsTable(), discardsSchema())
	if err != nil {
		switch {
		case errors.Is(err, errAuthz):
			m.setBackOff(err)
			validationError := m.validateConfig(ctx, asyncDest.Destination)
			if validationError != nil {
				err = fmt.Errorf("failed to validate snowpipe credentials: %s", validationError.Error())
			}
			return m.failedJobs(asyncDest, err.Error())
		case errors.Is(err, errBackoff):
			return m.failedJobs(asyncDest, err.Error())
		default:
			return m.abortJobs(asyncDest, fmt.Errorf("failed to prepare discards channel: %w", err).Error())
		}
	}
	m.logger.Infon("Prepared discards channel")

	groupedEvents := lo.GroupBy(events, func(event *event) string {
		return event.Message.Metadata.Table
	})
	uploadInfos := lo.MapToSlice(groupedEvents, func(tableName string, tableEvents []*event) *uploadInfo {
		jobIDs := lo.Map(tableEvents, func(event *event, _ int) int64 {
			return event.Metadata.JobID
		})
		latestJobID := lo.MaxBy(tableEvents, func(a, b *event) bool {
			return a.Metadata.JobID > b.Metadata.JobID
		})
		return &uploadInfo{
			tableName:              tableName,
			events:                 tableEvents,
			jobIDs:                 jobIDs,
			eventsSchema:           schemaFromEvents(tableEvents),
			discardChannelResponse: discardsChannel,
			latestJobID:            latestJobID.Metadata.JobID,
		}
	})
	slices.SortFunc(uploadInfos, func(a, b *uploadInfo) int {
		return int(a.latestJobID - b.latestJobID)
	})

	var (
		importingJobIDs, failedJobIDs []int64
		importInfos                   []*importInfo
		discardImportInfo             *importInfo
		failedReason                  string
	)
	shouldResetBackoff := true // backoff should be reset if authz error is not encountered for any of the tables
	isBackoffSet := false      // should not be set again if already set
	for _, info := range uploadInfos {
		imInfo, discardImInfo, err := m.sendEventsToSnowpipe(ctx, asyncDest.Destination.ID, &destConf, info)
		if err != nil {
			switch {
			case errors.Is(err, errAuthz):
				shouldResetBackoff = false
				if !isBackoffSet {
					isBackoffSet = true
					m.setBackOff(err)
					validationError := m.validateConfig(ctx, asyncDest.Destination)
					if validationError != nil && failedReason == "" {
						failedReason = fmt.Sprintf("failed to validate snowpipe credentials: %s", validationError.Error())
					}
				}
			case errors.Is(err, errBackoff):
				shouldResetBackoff = false
			}
			m.logger.Warnn("Failed to send events to Snowpipe",
				logger.NewStringField("table", info.tableName),
				obskit.Error(err),
			)
			if failedReason == "" {
				failedReason = err.Error()
			}
			failedJobIDs = append(failedJobIDs, info.jobIDs...)
			continue
		}

		importingJobIDs = append(importingJobIDs, info.jobIDs...)
		importInfos = append(importInfos, imInfo)

		if discardImInfo != nil && discardImportInfo == nil {
			discardImportInfo = discardImInfo
		} else if discardImInfo != nil {
			discardImportInfo.Count += discardImInfo.Count
			discardImportInfo.Offset = discardImInfo.Offset
		}
	}
	if shouldResetBackoff {
		m.resetBackoff()
	}
	if discardImportInfo != nil {
		importInfos = append(importInfos, discardImportInfo)
	}

	var importParameters stdjson.RawMessage
	if len(importInfos) > 0 {
		importIDBytes, err := jsonrs.Marshal(importInfos)
		if err != nil {
			return m.abortJobs(asyncDest, fmt.Errorf("failed to marshal import id: %w", err).Error())
		}

		importParameters = stdjson.RawMessage(`{"importId":` + string(importIDBytes) + `}`)
	}
	m.logger.Infon("Uploaded data to snowpipe streaming destination")

	m.stats.jobs.importing.Count(len(importingJobIDs))
	m.stats.jobs.failed.Count(len(failedJobIDs))

	return common.AsyncUploadOutput{
		ImportingJobIDs:     importingJobIDs,
		ImportingCount:      len(importingJobIDs),
		ImportingParameters: importParameters,
		FailedJobIDs:        failedJobIDs,
		FailedReason:        failedReason,
		FailedCount:         len(failedJobIDs),
		DestinationID:       asyncDest.Destination.ID,
	}
}

func (m *Manager) eventsFromFile(fileName string, eventsCount int) ([]*event, error) {
	file, err := os.Open(fileName)
	if err != nil {
		return nil, fmt.Errorf("failed to open file %s: %w", fileName, err)
	}
	defer func() {
		_ = file.Close()
	}()

	events := make([]*event, 0, eventsCount)

	formattedTS := m.Now().Format(misc.RFC3339Milli)
	scanner := bufio.NewScanner(file)
	scanner.Buffer(nil, int(m.config.maxBufferCapacity.Load()))

	for scanner.Scan() {
		var e event
		if err := jsonrs.Unmarshal(scanner.Bytes(), &e); err != nil {
			return nil, fmt.Errorf("failed to unmarshal event: %w", err)
		}
		e.setUUIDTimestamp(formattedTS)

		events = append(events, &e)
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("error reading from file: %w", err)
	}
	return events, nil
}

// sendEventsToSnowpipe sends events to Snowpipe for the given table.
// It creates a channel for the table, inserts the events into the channel, and sends the discard events to the discards table.
// It returns the import info for the table and the discard import info if any.
// In case of failure, it deletes the channel.
func (m *Manager) sendEventsToSnowpipe(
	ctx context.Context,
	destinationID string,
	destConf *destConfig,
	info *uploadInfo,
) (*importInfo, *importInfo, error) {
	offset := strconv.FormatInt(info.latestJobID, 10)

	log := m.logger.Withn(
		logger.NewStringField("table", info.tableName),
		logger.NewIntField("events", int64(len(info.events))),
		logger.NewStringField("offset", offset),
	)
	log.Infon("Sending events to Snowpipe")

	channelResponse, err := m.initializeChannelWithSchema(ctx, destinationID, destConf, info.tableName, info.eventsSchema)
	if err != nil {
		return nil, nil, fmt.Errorf("creating channel %s: %w", info.tableName, err)
	}
	log.Infon("Prepared channel", logger.NewStringField("channelID", channelResponse.ChannelID))

	formattedTS := m.Now().Format(misc.RFC3339Milli)
	var discardInfos []discardInfo
	for _, tableEvent := range info.events {
		discardInfos = append(discardInfos, getDiscardedRecordsFromEvent(m.logger, tableEvent, channelResponse.SnowpipeSchema, info.tableName, formattedTS)...)
	}

	/*
		Discards table events are inserted before main table events.
		This is because if discards insert succeeds but main table insert fails,
		then on retry, duplicate records will occur in the discards table.
		But if events are inserted in the reverse order, then on retry,
		duplicate records will occur in the main table.
	*/
	var discardImInfo *importInfo
	if len(discardInfos) > 0 {
		log.Infon("Inserting events into discards table", logger.NewIntField("count", int64(len(discardInfos))))
		discardImInfo, err = m.sendDiscardEventsToSnowpipe(ctx, offset, info.discardChannelResponse.ChannelID, discardInfos)
		if err != nil {
			return nil, nil, fmt.Errorf("sending discard events to Snowpipe: %w", err)
		}
	}

	insertReq := &model.InsertRequest{
		Rows: lo.Map(info.events, func(event *event, _ int) model.Row {
			return event.Message.Data
		}),
		Offset: offset,
	}
	channelID, err := m.insert(ctx, destinationID, destConf, info, insertReq, channelResponse.ChannelID)
	if err != nil {
		return nil, nil, fmt.Errorf("inserting events: %w", err)
	}

	log.Infon("Sent events to Snowpipe")

	imInfo := &importInfo{
		ChannelID: channelID,
		Offset:    offset,
		Table:     info.tableName,
		Count:     len(info.events),
	}
	return imInfo, discardImInfo, nil
}

func (m *Manager) insert(ctx context.Context, destinationID string, destConf *destConfig, info *uploadInfo, insertReq *model.InsertRequest, channelID string) (string, error) {
	deleteChannel := func(tableName, chID string) {
		if err := m.deleteChannel(ctx, tableName, chID); err != nil {
			m.logger.Warnn("Failed to delete channel", obskit.Error(err))
		}
	}
	extractError := func(res *model.InsertResponse, tableName string) error {
		errorMessages := lo.Map(res.Errors, func(ie model.InsertError, _ int) string {
			return ie.Message
		})
		return fmt.Errorf("inserting data %s failed: %s %v", tableName, res.Code, errorMessages)
	}

	insertRes, err := m.api.Insert(ctx, channelID, insertReq)
	if err == nil {
		if !insertRes.Success {
			deleteChannel(info.tableName, channelID)
			return "", extractError(insertRes, info.tableName)
		}
		return channelID, nil
	}
	/*
		404 can happen in the case where the channel is cached in the rudder-server,
		but then the snowpipe service restarts making the channel id invalid.
		But we don't want to pre-emptively check for validity of the channel before every insert
		because such 404 errors are expected to be rare.
	*/
	if errors.Is(err, snowpipeapi.ErrChannelNotFound) {
		deleteChannel(info.tableName, channelID)
		m.logger.Infon("Channel not found, recreating channel", logger.NewStringField("channelID", channelID))
		var insertRes2 *model.InsertResponse
		recreatedChannel, err2 := m.initializeChannelWithSchema(ctx, destinationID, destConf, info.tableName, info.eventsSchema)
		defer func() {
			if err2 != nil || !insertRes2.Success {
				deleteChannel(info.tableName, recreatedChannel.ChannelID)
			}
		}()
		if err2 != nil {
			return "", fmt.Errorf("creating channel %s: %w", info.tableName, err2)
		}
		m.logger.Infon("Recreated channel", logger.NewStringField("channelID", recreatedChannel.ChannelID))

		insertRes2, err2 = m.api.Insert(ctx, recreatedChannel.ChannelID, insertReq)
		if err2 != nil {
			return "", fmt.Errorf("inserting data %s: %w", info.tableName, err2)
		}
		if !insertRes2.Success {
			return "", extractError(insertRes2, info.tableName)
		}
		m.logger.Infon("Inserted data into recreated channel")
		return recreatedChannel.ChannelID, nil
	}
	// For all other errors, clean up and return
	deleteChannel(info.tableName, channelID)
	return "", fmt.Errorf("inserting data %s: %w", info.tableName, err)
}

// schemaFromEvents builds a schema by iterating over events and merging their columns
// using a first-encountered type basis for each column.
func schemaFromEvents(events []*event) whutils.ModelTableSchema {
	columnsMap := make(whutils.ModelTableSchema)
	for _, e := range events {
		for col, typ := range e.Message.Metadata.Columns {
			if _, exists := columnsMap[col]; !exists {
				columnsMap[col] = typ
			}
		}
	}
	return columnsMap
}

func (m *Manager) abortJobs(asyncDest *common.AsyncDestinationStruct, abortReason string) common.AsyncUploadOutput {
	m.stats.jobs.aborted.Count(len(asyncDest.ImportingJobIDs))
	return common.AsyncUploadOutput{
		AbortJobIDs:   asyncDest.ImportingJobIDs,
		AbortCount:    len(asyncDest.ImportingJobIDs),
		AbortReason:   abortReason,
		DestinationID: asyncDest.Destination.ID,
	}
}

func (m *Manager) failedJobs(asyncDest *common.AsyncDestinationStruct, failedReason string) common.AsyncUploadOutput {
	m.stats.jobs.failed.Count(len(asyncDest.ImportingJobIDs))
	return common.AsyncUploadOutput{
		FailedJobIDs:  asyncDest.ImportingJobIDs,
		FailedCount:   len(asyncDest.ImportingJobIDs),
		FailedReason:  failedReason,
		DestinationID: asyncDest.Destination.ID,
	}
}

// Poll checks the status of multiple imports using the import ID from pollInput.
// For the once which have reached the terminal state (success or failure), it caches the import infos in polledImportInfoMap. Later if Poll is called again, it does not need to do the status check again.
// Once all the imports have reached the terminal state, if any imports have failed, it deletes the channels for those imports.
// It returns a PollStatusResponse indicating if any imports are still in progress or if any have failed or succeeded
func (m *Manager) Poll(pollInput common.AsyncPoll) common.PollStatusResponse {
	m.logger.Infon("Polling started")

	var importInfos []*importInfo
	err := jsonrs.Unmarshal([]byte(pollInput.ImportId), &importInfos)
	if err != nil {
		return common.PollStatusResponse{
			InProgress: false,
			StatusCode: http.StatusBadRequest,
			Complete:   true,
			HasFailed:  true,
			Error:      fmt.Errorf("failed to unmarshal import id: %w", err).Error(),
		}
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if anyInProgress := m.processPollImportInfos(ctx, importInfos); anyInProgress {
		m.stats.pollingInProgress.Increment()

		// Check for stuck pipeline batch
		now := m.Now()
		if m.pollingStartTime.IsZero() {
			m.pollingStartTime = now
		} else {
			duration := now.Sub(m.pollingStartTime)
			threshold := m.config.stuckPipelineThreshold.Load()
			if duration > threshold {
				m.logger.Warnn("Stuck snowpipe pipeline detected",
					logger.NewDurationField("duration", duration),
					logger.NewDurationField("threshold", threshold),
					logger.NewStringField("importID", pollInput.ImportId),
				)
			}
		}
		return common.PollStatusResponse{InProgress: true}
	}

	updatedImportInfos := lo.Map(importInfos, func(item *importInfo, index int) *importInfo {
		return m.polledImportInfoMap[item.ChannelID]
	})
	failedImports := lo.Filter(updatedImportInfos, func(info *importInfo, index int) bool {
		return info.Failed
	})
	m.cleanupFailedImports(ctx, failedImports)
	m.updateJobStatistics(updatedImportInfos)
	// Reset batch polling start time since all imports completed
	m.pollingStartTime = time.Time{}
	m.polledImportInfoMap = make(map[string]*importInfo)

	return m.buildPollStatusResponse(updatedImportInfos, failedImports)
}

func (m *Manager) processPollImportInfos(ctx context.Context, infos []*importInfo) bool {
	var anyInProgress bool
	for i := range infos {
		info := infos[i]

		if _, alreadyProcessed := m.polledImportInfoMap[info.ChannelID]; alreadyProcessed {
			continue
		}
		info.FailedJobIds = nil
		info.Failed = false

		inProgress, err := m.getImportStatus(ctx, info)
		if err != nil {
			info.Failed = true
			info.Reason = err.Error()
			m.polledImportInfoMap[info.ChannelID] = info
			continue
		}
		if !inProgress {
			m.polledImportInfoMap[info.ChannelID] = info
		}

		anyInProgress = anyInProgress || inProgress
	}
	return anyInProgress
}

func (m *Manager) getImportStatus(ctx context.Context, info *importInfo) (bool, error) {
	log := m.logger.Withn(
		logger.NewStringField("channelId", info.ChannelID),
		logger.NewStringField("expectedOffset", info.Offset),
		logger.NewStringField("table", info.Table),
	)
	log.Infon("Polling for import info")

	statusRes, err := m.api.GetStatus(ctx, info.ChannelID)
	if err == nil {
		log.Infon("Polled import info",
			logger.NewBoolField("success", statusRes.Success),
			logger.NewStringField("latestCommittedOffset", statusRes.Offset),
			logger.NewStringField("latestInsertedOffset", statusRes.LatestInsertedOffset),
			logger.NewBoolField("valid", statusRes.Valid),
		)
		if !statusRes.Valid || !statusRes.Success {
			return false, fmt.Errorf("invalid status response with valid: %t, success: %t", statusRes.Valid, statusRes.Success)
		}
		return isInProgress(statusRes, info, log)
	}
	/*
		404 can happen in the case where the channel is cached in the rudder-server,
		but then the snowpipe service restarts making the channel id invalid.
		During polling, we try to recreate the channel to check the status.
	*/
	if errors.Is(err, snowpipeapi.ErrChannelNotFound) {
		log.Infon("Channel not found during polling, attempting to recreate channel", logger.NewStringField("channelID", info.ChannelID))

		var destConf destConfig
		if decodeErr := destConf.Decode(m.destination.Config); decodeErr != nil {
			return false, fmt.Errorf("failed to decode destination config during channel recreation: %w", decodeErr)
		}

		req := buildCreateChannelRequest(m.destination.ID, m.config.instanceID, &destConf, info.Table)
		recreatedChannel, recreateErr := m.api.CreateChannel(ctx, req)
		if recreateErr != nil {
			return false, fmt.Errorf("recreating channel: %w", recreateErr)
		}
		m.channelCache.Store(info.Table, recreatedChannel)

		log.Infon("Recreated channel for polling", logger.NewStringField("channelID", recreatedChannel.ChannelID))

		// The new channel id is not being associated with the import info. This is because even if we do that,
		// it will update the in memory map and not the database
		// So the next time we poll, we will get the old channel id and not the new one

		statusRes2, err2 := m.api.GetStatus(ctx, recreatedChannel.ChannelID)
		if err2 != nil {
			return false, fmt.Errorf("getting status after channel recreation: %w", err2)
		}

		log.Infon("Polled import info after recreation",
			logger.NewBoolField("success", statusRes2.Success),
			logger.NewStringField("latestCommittedOffset", statusRes2.Offset),
			logger.NewStringField("latestInsertedOffset", statusRes2.LatestInsertedOffset),
			logger.NewBoolField("valid", statusRes2.Valid),
		)

		if !statusRes2.Valid || !statusRes2.Success {
			return false, fmt.Errorf("invalid status response after recreation with valid: %t, success: %t", statusRes2.Valid, statusRes2.Success)
		}

		return isInProgress(statusRes2, info, log)
	}
	return false, fmt.Errorf("getting status: %w", err)
}

func isInProgress(statusRes *model.StatusResponse, info *importInfo, log logger.Logger) (bool, error) {
	latestCommittedOffset, err := convertToInt(statusRes.Offset)
	if err != nil {
		return false, fmt.Errorf("failed to convert latestCommittedOffset to int: %w", err)
	}
	latestInsertedOffset, err := convertToInt(statusRes.LatestInsertedOffset)
	if err != nil {
		return false, fmt.Errorf("failed to convert latestInsertedOffset to int: %w", err)
	}
	expectedOffset, err := convertToInt(info.Offset)
	if err != nil {
		return false, fmt.Errorf("failed to convert expectedOffset to int: %w", err)
	}

	// Case 1: All events have been flushed - proceed to next batch
	if latestCommittedOffset == expectedOffset {
		log.Infon("All events have been flushed successfully")
		return false, nil
	}

	// Case 2: Events lost - restart/error scenario
	if latestInsertedOffset < expectedOffset {
		log.Infon("Events lost due to Snowpipe restart or error")
		info.FailedJobIds = &failedJobIds{
			Start: latestCommittedOffset + 1,
			End:   expectedOffset,
		}
		return false, fmt.Errorf("events lost: latestCommittedOffset=%d, latestInsertedOffset=%d, expectedOffset=%d",
			latestCommittedOffset, latestInsertedOffset, expectedOffset)
	}

	// Case 3: Flushing in progress - continue polling
	if latestInsertedOffset > latestCommittedOffset {
		log.Debugn("Flushing in progress, continuing to poll")
		return true, nil
	}

	// Unexpected case - should not reach here based on the logic
	log.Warnn("Unexpected polling state encountered")
	return true, nil
}

func convertToInt(a string) (int64, error) {
	if a == "" {
		// Using math.MinInt64 instead of 0 to avoid skipping events
		// with negative job ids in case of downscaling of nodes
		return math.MinInt64, nil
	}
	ai, err := strconv.ParseInt(a, 10, 64)
	if err != nil {
		return 0, fmt.Errorf("failed to convert a to int: %w", err)
	}
	return ai, nil
}

func (m *Manager) cleanupFailedImports(ctx context.Context, failedInfos []*importInfo) {
	for _, info := range failedInfos {
		m.logger.Warnn("Failed to poll channel offset",
			logger.NewStringField("channelId", info.ChannelID),
			logger.NewStringField("offset", info.Offset),
			logger.NewStringField("table", info.Table),
			logger.NewStringField("reason", info.Reason),
		)

		if err := m.deleteChannel(ctx, info.Table, info.ChannelID); err != nil {
			m.logger.Warnn("Failed to delete channel",
				logger.NewStringField("channelId", info.ChannelID),
				logger.NewStringField("table", info.Table),
				obskit.Error(err),
			)
		}
	}
}

func (m *Manager) updateJobStatistics(importInfos []*importInfo) {
	var successfulCount, failedCount int

	for _, info := range importInfos {
		if info.Failed {
			failedCount += info.Count
		} else {
			successfulCount += info.Count
		}
	}
	m.stats.jobs.failed.Count(failedCount)
	m.stats.jobs.succeeded.Count(successfulCount)
}

func (m *Manager) buildPollStatusResponse(importInfos, failedImports []*importInfo) common.PollStatusResponse {
	response := common.PollStatusResponse{
		InProgress: false,
		StatusCode: http.StatusOK,
		Complete:   true,
	}

	if len(failedImports) > 0 {
		response.HasFailed = true
		response.FailedJobParameters = stringify.Any(importInfos)
	} else {
		response.HasFailed = false
		response.HasWarning = false
	}
	return response
}

// GetUploadStats returns the status of the uploads for the snowpipe streaming destination.
// It returns the status of the uploads for the given job IDs.
// If any of the uploads have failed, it returns the reason for the failure.
func (m *Manager) GetUploadStats(input common.GetUploadStatsInput) common.GetUploadStatsResponse {
	m.logger.Infon("Getting import stats for snowpipe streaming destination")

	var importInfos []*importInfo
	err := jsonrs.Unmarshal([]byte(input.FailedJobParameters), &importInfos)
	if err != nil {
		return common.GetUploadStatsResponse{
			StatusCode: http.StatusBadRequest,
			Error:      fmt.Errorf("failed to unmarshal failed job urls: %w", err).Error(),
		}
	}

	succeededTables, failedTables := make(map[string]struct{}), make(map[string]*importInfo)
	for _, info := range importInfos {
		if info.Failed {
			failedTables[info.Table] = info
		} else {
			succeededTables[info.Table] = struct{}{}
		}
	}

	var (
		succeededJobIDs  []int64
		failedJobIDs     []int64
		failedJobReasons = make(map[int64]string)
	)
	for _, job := range input.ImportingList {
		tableName := gjson.GetBytes(job.EventPayload, "metadata.table").String()
		if _, ok := succeededTables[tableName]; ok {
			succeededJobIDs = append(succeededJobIDs, job.JobID)
		}
		if info, ok := failedTables[tableName]; ok {
			if info.FailedJobIds == nil || (job.JobID >= info.FailedJobIds.Start && job.JobID <= info.FailedJobIds.End) {
				failedJobIDs = append(failedJobIDs, job.JobID)
				failedJobReasons[job.JobID] = info.Reason
			} else {
				succeededJobIDs = append(succeededJobIDs, job.JobID)
			}
		}

	}
	return common.GetUploadStatsResponse{
		StatusCode: http.StatusOK,
		Metadata: common.EventStatMeta{
			FailedKeys:    failedJobIDs,
			FailedReasons: failedJobReasons,
			SucceededKeys: succeededJobIDs,
		},
	}
}

func (m *Manager) isInBackoff() bool {
	if m.backoff.next.IsZero() {
		return false
	}
	return m.Now().Before(m.backoff.next)
}

func (m *Manager) resetBackoff() {
	m.backoff.next = time.Time{}
	m.backoff.attempts = 0
	m.backoff.error = ""
}

func (m *Manager) setBackOff(err error) {
	b := backoff.NewExponentialBackOff(
		backoff.WithInitialInterval(m.config.backoff.initialInterval.Load()),
		backoff.WithMultiplier(m.config.backoff.multiplier.Load()),
		backoff.WithClockProvider(m),
		backoff.WithRandomizationFactor(0),
		backoff.WithMaxElapsedTime(0),
		backoff.WithMaxInterval(m.config.backoff.maxInterval.Load()),
	)
	b.Reset()
	m.backoff.attempts++
	m.backoff.error = err.Error()

	var d time.Duration
	for index := int64(0); index < int64(m.backoff.attempts); index++ {
		d = b.NextBackOff()
	}
	m.backoff.next = m.Now().Add(d)
}

func buildCreateChannelRequest(destinationID, partition string, destConf *destConfig, tableName string) *model.CreateChannelRequest {
	return &model.CreateChannelRequest{
		RudderIdentifier: destinationID,
		Partition:        partition,
		AccountConfig: model.AccountConfig{
			Account:              destConf.Account,
			User:                 destConf.User,
			Role:                 destConf.Role,
			PrivateKey:           whutils.FormatPemContent(destConf.PrivateKey),
			PrivateKeyPassphrase: destConf.PrivateKeyPassphrase,
		},
		TableConfig: model.TableConfig{
			Database:      destConf.Database,
			Schema:        destConf.Namespace,
			Table:         tableName,
			EnableIceberg: destConf.EnableIceberg,
		},
	}
}
